import os
import pandas as pd
import numpy as np
import logging
import time
import akshare as ak
from qtorch.core.interfaces import IDataProvider
from qtorch.core.datafetcher import DataFetcher
from qtorch.core.exceptions import FinancialDataException
from qtorch.utils.validators import ValidationError
from typing import Optional
from qtorch.tool import retry

logging.basicConfig(level=logging.INFO)

column_mapping = {
    '日期': 'date',
    '开盘': 'open',
    '收盘': 'close',
    '最高': 'high',
    '最低': 'low',
    '成交量': 'volume'
}

class AKShareStockFetcher(DataFetcher, IDataProvider):
    def __init__(self, cache_dir="data_cache", cache_ttl=86400):
        super().__init__()
        self.symbol_map = {
            "sh510300": "510300.SH",  # 沪深300ETF
            "sz159919": "159919.SZ",  # 沪深300ETF(深)
            "IF888.CFX": "沪深300指数期货",
            "RB2305.SHF": "螺纹钢期货"
        }
        self.cache_dir = cache_dir
        self.cache_ttl = cache_ttl
        os.makedirs(self.cache_dir, exist_ok=True)

    def _get_cache_path(self, symbol: str) -> str:
        return os.path.join(self.cache_dir, f"{symbol}.parquet")

    def _load_from_cache(self, symbol: str) -> Optional[pd.DataFrame]:
        cache_path = self._get_cache_path(symbol)
        if os.path.exists(cache_path):
            file_mtime = os.path.getmtime(cache_path)
            if time.time() - file_mtime < self.cache_ttl:
                try:
                    return pd.read_parquet(cache_path)
                except Exception as e:
                    logging.warning(f"缓存加载失败: {e}")
        return None

    def _save_to_cache(self, symbol: str, data: pd.DataFrame):
        try:
            cache_path = self._get_cache_path(symbol)
            data.to_parquet(cache_path)
        except Exception as e:
            logging.error(f"缓存保存失败: {e}")

    @retry(max_retries=3, delay=5)
    def fetch_financials(self, symbol: str, report_type: str = 'annual') -> pd.DataFrame:
        """实现AKShare财务数据获取"""
        try:
            logging.info(f"正在获取 {symbol} {report_type}财务数据...")
            financial_df = ak.stock_financial_report(symbol,
                report_type="年报" if report_type == 'annual' else "季报")
            
            # 财务数据字段标准化
            financial_df.rename(columns={
                '报告日期': 'report_date',
                '每股收益': 'eps',
                '资产负债率': 'debt_ratio',
                '流动比率': 'current_ratio',
                '净利润': 'net_profit',
                '营业收入': 'revenue'
            }, inplace=True)
            
            # 类型转换和日期处理
            financial_df['report_date'] = pd.to_datetime(financial_df['report_date'])
            financial_df.set_index('report_date', inplace=True)
            
            # 关键指标验证
            required_columns = ['eps', 'debt_ratio', 'current_ratio']
            missing_cols = [col for col in required_columns if col not in financial_df.columns]
            if missing_cols:
                raise FinancialDataException(f"财务数据缺失关键字段: {missing_cols}")
                
            return financial_df[required_columns + ['net_profit', 'revenue']]
            
        except Exception as e:
            logging.error(f"财务数据获取失败: {symbol} - {str(e)}")
            raise FinancialDataException(f"财务数据获取异常: {symbol}") from e

    @retry(max_retries=3, delay=5)

    def get_market_data(self, symbol: str, start_date: str, end_date: str) -> pd.DataFrame:
        """获取市场数据"""
        try:
            # 检查日期格式
            pd.to_datetime(start_date)
            pd.to_datetime(end_date)
            
            # 获取数据
            data = self.fetch_data(symbol, start_date, end_date)
            
            # 数据验证
            if data.empty:
                logging.warning(f"未获取到{symbol}的有效数据")
                return pd.DataFrame()
                
            # 数据预处理
            processed_data = self.preprocess_data(data)
            
            # 添加市场标识
            processed_data['is_market'] = True
            
            return processed_data
            
        except Exception as e:
            logging.error(f"获取{symbol}数据失败: {str(e)}")
            raise FinancialDataException(f"市场数据获取失败: {symbol}") from e
        
    def _fetch_data_impl(self, symbol: str, start_date: str, end_date: str) -> pd.DataFrame:
        """实际数据获取实现"""
        try:
            logging.info(f"正在获取 {symbol} 数据...")
            # 优先从缓存加载
            cached_data = self._load_from_cache(symbol)
            if cached_data is not None:
                return cached_data[(cached_data.index >= start_date) & (cached_data.index <= end_date)]
                
            # 根据不同类型获取数据
            if symbol.endswith('.CFX'):
                logging.warning(f"期货数据获取暂不支持: {symbol}")
                return pd.DataFrame()
            else:
                df = ak.stock_zh_a_hist(symbol=symbol, period="daily")
            # 重命名列
            df.rename(columns=column_mapping, inplace=True)
            # 转换日期格式
            df['date'] = pd.to_datetime(df['date'])
            df.set_index('date', inplace=True)
            # 选择所需列
            df = df[['open', 'high', 'low', 'close', 'volume']]
            # 转换日期格式
            start_dt = pd.to_datetime(start_date)
            end_dt = pd.to_datetime(end_date)
            # 过滤日期范围
            df = df[(df.index >= start_dt) & (df.index <= end_dt)]
            if df.empty:
                logging.warning(f"在 {start_date} 到 {end_date} 范围内未找到 {symbol} 的数据")
                return df
                
            # 数据验证
            try:
                # 确保所有列都是数值类型
                df = df.apply(pd.to_numeric, errors='coerce')
                # 检查是否有缺失值
                if df.isnull().values.any():
                    raise ValueError("数据包含无效值")
                # 检查是否有数据
                if len(df) == 0:
                    raise ValueError("未找到指定日期范围内的数据")
            except Exception as e:
                logging.error(f"数据验证失败: {str(e)}")
                raise ValueError("数据完整性检查未通过") from e

            self._save_to_cache(symbol, df)
            return df
        except Exception as e:
            logging.error(f"数据获取失败: {e}")
            return pd.DataFrame()

    @retry(max_retries=3, delay=5)
    def fetch_benchmark_data(self, benchmark_code: str, start_date: str, end_date: str) -> pd.DataFrame:
        """获取基准指数数据（如沪深300）"""
        try:
            # 增强缓存键唯一性（增加基准类型前缀）
            cache_key = f"benchmark_{benchmark_code}_{start_date}_{end_date}"
            cached_data = self._load_from_cache(cache_key)
            
            if cached_data is not None:
                logging.info(f"从缓存加载基准数据 {benchmark_code} [{start_date} 至 {end_date}]")
                return cached_data

            logging.info(f"正在获取基准数据 {benchmark_code} [{start_date} 至 {end_date}]...")
            
            # 根据不同类型获取基准数据
            if benchmark_code.startswith('sh') or benchmark_code.startswith('sz'):
                df = ak.stock_zh_index_daily(symbol=benchmark_code)
            else:
                df = ak.fund_etf_hist_sina(symbol=benchmark_code)
            
            # 统一字段处理
            df.rename(columns=column_mapping, inplace=True)
            
            # 统一日期处理（与股票数据处理逻辑一致）
            df.index = pd.to_datetime(df["date"])
            start_idx = df.index.get_loc(start_date)
            end_idx = df.index.get_loc(end_date) + 1
            df = df.sort_index().iloc[start_idx:end_idx]
            
            # 填充缺失日期（增加节假日处理）
            df = self._fill_missing_dates(df, start_date, end_date)
            # 添加基准标识
            df['is_benchmark'] = True
            
            if df.empty:
                raise ValueError(f"未找到 {benchmark_code} 的基准数据")
                
            # 缓存处理前验证数据
            from qtorch.utils.validators import (
                validate_dataframe_structure,
                validate_timestamp_format,
                ValidationError
            )
            
            required_columns = ['open', 'high', 'low', 'close', 'volume']
            try:
                validate_dataframe_structure(df, required_columns)
                validate_timestamp_format(df)
            except ValidationError as e:
                logging.error(f"基准数据验证失败: {e}")
                raise ValueError("基准数据完整性检查未通过") from e

            self._save_to_cache(cache_key, df)
            return self.preprocess_data(df[['open', 'high', 'low', 'close', 'volume']])
            
        except Exception as e:
            logging.error(f"基准数据获取失败: {e}")
            return pd.DataFrame()

    def _fill_missing_dates(self, df: pd.DataFrame, start: str, end: str) -> pd.DataFrame:
        """填充缺失的交易日数据"""
        full_index = pd.date_range(start=start, end=end, freq='B')
        return df.reindex(full_index).ffill().dropna()

    def _align_dates(self, df: pd.DataFrame, start: str, end: str) -> pd.DataFrame:
        """确保数据日期范围精确匹配"""
        try:
            # 转换日期格式
            start_date = pd.to_datetime(start)
            end_date = pd.to_datetime(end)
            
            # 创建数据副本
            aligned_df = df.copy()
            
            # 获取日期索引位置
            try:
                # 使用get_indexer确保返回整数索引
                start_idx, end_idx = aligned_df.index.get_indexer([start_date, end_date])
                end_idx += 1
            except KeyError as e:
                raise ValueError(f"日期范围不匹配: {str(e)}")
            # 验证索引范围
            if start_idx < 0 or end_idx > len(aligned_df):
                raise ValueError("无效的日期范围")
            # 执行切片
            aligned_df = aligned_df.iloc[start_idx:end_idx]
            # 验证结果
            if aligned_df.empty:
                raise ValueError("未找到指定日期范围内的数据")
            return aligned_df
        except Exception as e:
            logging.error(f"日期对齐失败: {str(e)}")
            raise FinancialDataException("日期对齐处理失败") from e
            logging.error(f"日期对齐失败: {str(e)}")
            raise FinancialDataException("日期对齐处理失败") from e

    def preprocess_data(self, raw_data: pd.DataFrame) -> pd.DataFrame:
        """统一预处理股票和基准数据"""
        df = raw_data.copy()
        # 统一计算收益率
        df["returns"] = df["close"].pct_change()
        # 计算波动率
        df["volatility"] = df["returns"].rolling(5).std()
        # 清理无效数据
        df.dropna(inplace=True)
        # 保留收盘价用于基准计算
        return df[["open", "high", "low", "close", "volume", "returns", "volatility"]]
